#!/usr/bin/env python

import torch
from setuptools import setup, find_packages, Extension
# from torch.utils import cpp_extension
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension

# setup(name='torchfwi',
#       ext_modules=[
#             cpp_extension.CppExtension('torchfwi', ['Torch_Fwi.cpp', 'Parameter.cpp', 'libCUFD.cu', 'el_stress.cu', 'el_velocity.cu', 'el_stress_adj.cu', 'el_velocity_adj.cu', 'Model.cu', 'Cpml.cu', 'utilities.cu',	'Src_Rec.cu', 'Boundary.cu'])],
#       cmdclass={'build_ext': cpp_extension.BuildExtension})

# module = Extension('torchfwi',
#       sources=['src/cuda_ext/utilities.cu'],
#       include_dirs=[r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\include'],
#       library_dirs=[r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\lib\x64'],
#       libraries=['cufft']
#   )

setup(name='fwi',
      package_dir={"": "src"},  # 关键配置
      packages=find_packages(where="src"),
      # package_data={
      #   "Torch_Fwi.cpp": ["*.pyd"]
      # },
      version='0.1.0',
      install_requires=[
          'torch',
          'numpy',
          'matplotlib'
      ],
      author='Dongzhuo Li, Kailai Xu, Jerry M. Harris, etc, modified by xbs',
      author_email='xbs150@163.com',
      description='Elastic Full-Waveform Inversion Integrated with PyTorch',
      long_description=open('README.md').read(),
      long_description_content_type='text/markdown',
      url='https://github.com/lidongzh/TorchFWI',
      classifiers=[
          'Programming Language :: Python :: 3',
          'License :: OSI Approved :: MIT License',
          'Operating System :: OS Independent',
      ],
      python_requires='>=3.9',
      ext_modules=[
          # CppExtension(name='torchfwi',
          #              sources=['src/cuda_ext/Torch_Fwi.cpp',
          #                       'src/cuda_ext/Parameter.cpp',
          #                       'src/cuda_ext/Boundary.cu',
          #                       'src/cuda_ext/libCUFD.cu',
          #                       'src/cuda_ext/el_stress.cu',
          #                       'src/cuda_ext/el_velocity.cu',
          #                       # 'src/cuda_ext/image_vel.cu',
          #                       # 'src/cuda_ext/image_vel_time.cu',
          #                       'src/cuda_ext/el_stress_adj.cu',
          #                       'src/cuda_ext/el_velocity_adj.cu',
          #                       'src/cuda_ext/Model.cu',
          #                       'src/cuda_ext/Cpml.cu',
          #                       'src/cuda_ext/utilities.cu',
          #                       'src/cuda_ext/Src_Rec.cu'
          #                       ],
          #              include_dirs=['src/cuda_ext/','src/cuda_ext/rapidjson', torch.utils.cpp_extension.include_paths()],
          #              extra_compile_args={'cxx': ['-std=c++14'], 'nvcc': ['-O2']}
          #              ),
          CUDAExtension(name='fwi_cuda_impl',
                        sources=['src/cuda_ext/Torch_Fwi.cpp',
                                 'src/cuda_ext/Parameter.cpp',
                                 'src/cuda_ext/libCUFD.cu',
                                 'src/cuda_ext/el_stress.cu',
                                 'src/cuda_ext/el_velocity.cu',
                                 'src/cuda_ext/el_stress_adj.cu',
                                 'src/cuda_ext/el_velocity_adj.cu',
                                 'src/cuda_ext/image_vel.cu',
                                 'src/cuda_ext/image_vel_time.cu',
                                 'src/cuda_ext/Model.cu',
                                 'src/cuda_ext/Cpml.cu',
                                 'src/cuda_ext/utilities.cu',
                                 'src/cuda_ext/Src_Rec.cu',
                                 'src/cuda_ext/Boundary.cu'],
                        # include_dirs=['src/cuda_ext/','src/cuda_ext/rapidjson',  torch.utils.cpp_extension.include_paths(),r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\include'],
                        include_dirs=['src/cuda_ext/', 'src/cuda_ext/rapidjson',
                                      torch.utils.cpp_extension.include_paths()],
                        # library_dirs=[r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\lib\x64'],
                        libraries=['cufft', 'nvrtc', 'cuda', 'cudart'],
                        extra_compile_args={'cxx': ['-std=c++17'], 'nvcc': ['-O2']}
                        ),
          # module
      ],
      cmdclass={
          'build_ext': BuildExtension
      }
      )
